package com.johnguo.beam_test.java;

import java.util.Arrays;
import java.util.List;

import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.options.PipelineOptions;
import org.apache.beam.sdk.options.PipelineOptionsFactory;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.DoFn;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Partition;
import org.apache.beam.sdk.transforms.Partition.PartitionFn;
import org.apache.beam.sdk.transforms.Flatten;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.sdk.values.PCollectionList;

public class PartitionDemo {
	static final List<String> PERSON = Arrays.asList("johnguo 24", "annasui 30", "liudan 25");

	public static void main(String... args) {
		PipelineOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().create();
		Pipeline p = Pipeline.create(options);
		PCollection<String> person = p.apply(Create.of(PERSON)).setCoder(StringUtf8Coder.of());
		PCollectionList<String> collectionList = person.apply(Partition.of(5, new PartitionFn<String>() {
			private static final long serialVersionUID = 1L;

			@Override
			public int partitionFor(String elem, int numPartitions) {
				String[] p = elem.split(" ");
				return Integer.valueOf(p[1]) % 5;
			}

		}));
		collectionList.get(0).apply(ParDo.of(new DoFn<String, String>() {

			private static final long serialVersionUID = 1L;

			@ProcessElement
			public void processElement(ProcessContext c) {
				System.out.println(c.element());
			}
		}));
		collectionList.apply(Flatten.<String>pCollections());
		p.run().waitUntilFinish();
	}
}
